# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import os
import json
import time
import logging
from events import scheduler
from common.cache.queue_cache import QueueCache
from common.enums.apply import MusicEnum
from common.enums.wallet import WalletEnum
from common.enums.public import CrontabEnum
from common.utils.cache import RedisUtil
from common.utils.times import TimeUtil
from common.utils.tools import ToolsUtil
from common.utils.urls import UrlUtil
from common.models.sys import SysCrontabModel
from common.models.users import UserWalletModel
from common.models.music import MusicRecordsModel
from common.chain.music_server import MusicOpenAiHkServer
from plugins.storage.driver import StorageDriver


logger = logging.getLogger(__name__)


async def execute(**kwargs):
    """ AI引擎回调任务 """
    start_time = time.time()
    crontab_id: int = int(kwargs["w_id"])
    crontab_job: str = str(kwargs["w_job"])
    if not scheduler.get_job(crontab_job):
        scheduler.pause_job(crontab_job)

    try:
        record_id = await QueueCache.pop_queue_value(QueueCache.MUSIC_JOB)
        if not record_id or record_id is None:
            return False

        await dispense(int(record_id))

        await SysCrontabModel.compute(crontab_id, start_time)
    except Exception as e:
        await SysCrontabModel.compute(crontab_id, start_time, status=CrontabEnum.CRON_ERROR, error=str(e))


async def dispense(record_id: int):
    # 查询任务
    record = await MusicRecordsModel.filter(id=record_id).first()
    if not record and record.status != MusicEnum.STATUS_ING:
        return

    # 处理响应
    if record.engine == "suno":
        await handle_suno_check(record)
    else:
        record.status = MusicEnum.STATUS_FAIL
        record.fail_reason = f"No processing channel {record.engine}"
        await record.save()


async def handle_suno_check(record):
    """
    Suno音乐回调处理

    Document:
        https://openai-hk.com/docs/lab/suno-v3.html

    Author:
        zero
    """
    cover_dir_path = ""
    audio_dir_path = ""
    response = {}
    try:
        # 限流拦截
        lock_key = f"music:prevent_{record.id}"
        lock_cache = await RedisUtil.get(lock_key)
        if lock_cache:
            await QueueCache.push_music(record.id)
            return False

        # 锁定记录
        await RedisUtil.set(lock_key, record.id, 60)
        print(f"Suno音乐开始: [重试={record.retry_count}, 任务={record.id}, 时间={time.time()}]")

        # 获取生成结果
        server = MusicOpenAiHkServer()
        response = await server.fetch(record.task_id)
        print(response)
        # 获取结果异常
        if isinstance(response, dict):
            error = response.get("error") or {}
            await MusicRecordsModel.filter(id=record.id).update(
                status=int(MusicEnum.STATUS_FAIL),
                fail_reason=error.get("message", "未知错误"),
                notify_snap=json.dumps(response, ensure_ascii=False),
                update_time=int(time.time())
            )
            await UserWalletModel.inc(
                scene=WalletEnum.POINTS,
                user_id=record.user_id,
                change_type=WalletEnum.POINTS_INC_VIDEO_FAIL,
                change_amount=record.use_points,
                project="AI音乐",
                remarks=error.get("message", "未知错误")
            )
            return None

        # 处理生成结果
        if response[0].get("status") == "complete":
            # 成功
            cover_url: str = response[0]["image_large_url"]
            audio_url: str = response[0]["audio_url"]
            duration: int = int(response[0]["metadata"].get("duration", 0))
            lyric: str = response[0]["metadata"].get("gpt_description_prompt", "")
            tags: str = response[0]["metadata"].get("tags", "")

            # 存储路径
            filename: str = ToolsUtil.make_md5_str(str(audio_url)+str(time.time()))
            curr_now: str = TimeUtil.timestamp_to_date(int(time.time()), str("%Y%m%d"))
            cover_dir: str = f"storage/applies/music/{curr_now}/{filename}.jpeg"
            audio_dir: str = f"storage/applies/music/{curr_now}/{filename}.mp3"
            cover_dir_path: str = UrlUtil.to_root_path(path=cover_dir)
            audio_dir_path: str = UrlUtil.to_root_path(path=audio_dir)

            # 下载封面
            await ToolsUtil.download_file(url=cover_url, save_path=cover_dir_path.strip())

            # 下载音频
            await ToolsUtil.download_file(url=audio_url, save_path=audio_dir_path.strip())

            # 视频上传OSS
            engine: str = await UrlUtil.get_storage_engine()
            if engine != "local":
                await StorageDriver.push(path=cover_dir_path, key=cover_dir.strip())
                await StorageDriver.push(path=audio_dir_path, key=audio_dir.strip())

            # 更新记录
            await MusicRecordsModel.filter(id=record.id).update(
                status=MusicEnum.STATUS_YES,
                cover_url=str(cover_dir).strip(),
                audio_url=str(audio_dir).strip(),
                duration=duration,
                lyric=lyric,
                tags=tags,
                fail_reason="",
                notify_snap=""
            )

            # 生成成功
            task_time = int(time.time()) - int(record.start_time)
            print(f"Suno音乐成功: [重试={record.retry_count}, 任务={record.id}, 耗时={task_time}s]")
        elif response[0].get("status"):
            # 进行中
            await QueueCache.push_music(record.id)

    except Exception as e:
        if record.retry_count < 2:
            print(f"Suno音乐重试: [重试={record.retry_count + 1}, 任务={record.id}, 错误={str(e)}]")
            record.retry_count += 1
            record.start_time = int(time.time())
            await record.save()
            await QueueCache.push_music(record.id)
        else:
            logger.error("Suno音乐回调失败: [%s]", str(e))
            await MusicRecordsModel.filter(id=record.id).update(
                status=MusicEnum.STATUS_FAIL,
                notify_snap=json.dumps(response, ensure_ascii=False),
                fail_reason=str(e),
                update_time=int(time.time())
            )
            if record.use_points > 0:
                await UserWalletModel.inc(
                    scene=WalletEnum.POINTS,
                    user_id=record.user_id,
                    change_type=WalletEnum.POINTS_INC_VIDEO_FAIL,
                    change_amount=record.use_points,
                    project="AI音乐",
                    remarks=str(e)
                )
    finally:
        engine: str = await UrlUtil.get_storage_engine()
        if engine != "local":
            if cover_dir_path and os.path.exists(cover_dir_path):
                os.unlink(cover_dir_path)
            if audio_dir_path and os.path.exists(audio_dir_path):
                os.unlink(audio_dir_path)
